import torch.nn as nn
import torch.nn.functional as F

class CNNMnist(nn.Module):
    """
    用于MNIST数据集的卷积神经网络模型。

    该模型包含两个卷积层和两个全连接层，适用于MNIST数据集的手写数字识别任务。
    """
    def __init__(self):
        """
        初始化模型的卷积层和全连接层。
        """
        super(CNNMnist, self).__init__()

        # 第一个卷积层，输入通道为1，输出通道为32，卷积核大小为5
        self.conv1 = nn.Conv2d(1, 32, kernel_size=5)
        # 第二个卷积层，输入通道为32，输出通道为64，卷积核大小为5
        self.conv2 = nn.Conv2d(32, 64, kernel_size=5)
        # 第一个全连接层，输入大小为1024，输出大小为256
        self.fc1 = nn.Linear(1024, 256)
        # 第二个全连接层，输入大小为256，输出大小为10
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        """
        定义前向传播过程。

        参数:
        x (Tensor): 输入的图像数据。

        返回:
        Tensor: 经过网络的前向传播后的输出。
        """
        # 通过第一个卷积层，使用ReLU激活函数和最大池化
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        # 通过第二个卷积层，使用ReLU激活函数和最大池化
        x = F.relu(F.max_pool2d(self.conv2(x), 2))
        # 将卷积层的输出展平，以便输入到全连接层
        x = x.view(-1, 1024)
        # 通过第一个全连接层，使用ReLU激活函数
        x = F.relu(self.fc1(x))
        # 通过第二个全连接层
        x = self.fc2(x)
        # 使用log_softmax函数进行输出，适用于分类任务
        return F.log_softmax(x, dim=1)

## 加入Dropout解决过拟合问题
class CNNDropoutMnist(nn.Module):
    """
    CNN模型用于MNIST数据集分类，包含两个卷积层、两个全连接层以及Dropout层。

    该模型的结构如下：
    - 第一个卷积层：输入通道为1，输出通道为32，卷积核大小为3x3，步幅为1。
    - 第二个卷积层：输入通道为32，输出通道为64，卷积核大小为3x3，步幅为1。
    - 最大池化层：池化窗口大小为2x2。
    - 第一个全连接层：输入大小为64*5*5，输出大小为256。
    - 第二个全连接层：输入大小为256，输出大小为10（对应10个类别）。
    - Dropout层：丢弃率为0.25，用于防止过拟合。
    """
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1)
        self.pool = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(64 * 5 * 5, 256)
        self.fc2 = nn.Linear(256, 10)
        self.dropout = nn.Dropout(0.25)

    def forward(self, x):
        """
        前向传播函数，定义了数据通过模型的流程。

        参数:
        x (torch.Tensor): 输入数据，形状为(batch_size, 1, 28, 28)。

        返回值:
        torch.Tensor: 模型的输出，形状为(batch_size, 10)，表示每个类别的得分。
        """
        # 第一层卷积 + ReLU激活 + 最大池化
        x = self.pool(nn.functional.relu(self.conv1(x)))
        # 第二层卷积 + ReLU激活 + 最大池化
        x = self.pool(nn.functional.relu(self.conv2(x)))
        # 将特征图展平为一维向量
        x = x.view(-1, 64 * 5 * 5)
        # 全连接层 + ReLU激活
        x = nn.functional.relu(self.fc1(x))
        # 应用Dropout
        x = self.dropout(x)
        # 输出层
        x = self.fc2(x)
        return x
